import gradio as gr
import numpy as np
from PIL import Image
from pinecone import Pinecone, ServerlessSpec

# 初始化Pinecone客户端
def initialize_pinecone(api_key):
    try:
        pinecone = Pinecone(api_key=api_key, environment='us-east-1')
        print("Pinecone client initialized successfully.")
        return pinecone
    except Exception as e:
        print(f"Error initializing Pinecone client: {e}")
        return None

# 管理索引：检查索引是否存在，如果不存在则创建新索引
def manage_index(pinecone, index_name):
    existing_indexes = pinecone.list_indexes()
    if any(index['name'] == index_name for index in existing_indexes): #索引存在
        print(f"索引 '{index_name}' 已存在。")
    else: #索引不存在
        print(f"索引 '{index_name}' 不存在。")
        print(f"正在创建新索引 '{index_name}'...")
        pinecone.create_index(
            name=index_name,
            dimension=64,  # MNIST 每个图像展平后是一个 64 维向量
            metric="euclidean",  # 使用欧氏距离
            spec=ServerlessSpec(cloud="aws", region="us-east-1")  # 服务器规格
        )
        print(f"索引 '{index_name}' 创建成功。")
    return pinecone.Index(index_name)

# 将图像转换为Pinecone索引可以接受的格式
def image_to_vector(image):
   
    # 调整图像大小为 8x8
    image = image.resize((8, 8), Image.Resampling.LANCZOS)

    # 将 PIL 图像转换为 NumPy 数组
    image_array1 = np.array(image)

    image_array2 = image_array1[:, :, 3]
    threshold1 = 30
    image_array2[image_array2 < threshold1] = 0
    #归一化
    image_array2 = (image_array2-image_array2.min())/(image_array2.max()-image_array2.min())*16
    image_array2 = np.round(image_array2)
    print(image_array2)
    # 展平图像
    image_array3 = image_array2.reshape((1, -1))

    # 展平图像
    image_vector = image_array3.tolist()
    return image_vector

# 使用Pinecone进行预测
def predict_digit(drawing, pinecone, index_name):
    try:
        if pinecone is None:
            return "Pinecone client not initialized."

        # 获取PIL 图像
        image = drawing['composite']

        # 将图像转换为向量
        image_vector = image_to_vector(image)

        # 执行查询
        index = manage_index(pinecone, index_name)
        if index is None:
            return "Index not found or not initialized."

        results = index.query(vector=image_vector, top_k=1, include_metadata=True)

        # 获取最接近的向量并返回其标签
        if results['matches']:
            return str(results['matches'][0]['metadata']['label'])
        else:
            return "No match found"
    except Exception as e:
        return f"Error during prediction: {e}"

def create_gradio_interface(api_key, index_name):
    pinecone = initialize_pinecone(api_key)

    inp = gr.Sketchpad(label="Draw a digit", type='pil')  # 创建 Sketchpad 组件并指定类型

    def wrapped_predict(drawing):
        return predict_digit(drawing, pinecone, index_name)

    iface = gr.Interface(
        fn=wrapped_predict,
        inputs=inp,
        outputs="text",
        live=False,
        title="Digit Recognition with Pinecone",
        description="Draw a digit and get it recognized using Pinecone."
    )
    return iface
# 启动Gradio界面

if __name__ == "__main__":
    api_key = "0f4bb63e-9728-43d4-89b8-562f96acb31b"  # 替换为你的Pinecone API密钥
    index_name = "mnist-index"  # 替换为你的Pinecone索引名称
   
    demo = create_gradio_interface(api_key, index_name)
    demo.launch()
    